#codeing=utf-8

"""
Candidates detection. Training and predecting process

date: 2017.03.29
"""

import numpy as np
import scipy.io as sio
import h5py
import sys
import time
import os
import sys, subprocess
from easydict import EasyDict as edict

from tools.train_net import train_process
from tools.dj_test_net_vol_luna_kaggle import test_process
from tools.cdd_use_lungmask import cdd_filter_lungmask
from tools.gen_candis_list import gen_list_5fold, gen_list_one_fold


def main():
    log_file_name = time.strftime('log/train_test_%Y-%m-%d_%H-%M-%S.txt', time.localtime(time.time()))
    sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', 0)
    tee = subprocess.Popen(["tee", log_file_name], stdin=subprocess.PIPE)
    os.dup2(tee.stdin.fileno(), sys.stdout.fileno())
    os.dup2(tee.stdin.fileno(), sys.stderr.fileno())

    cfg = edict()

    cfg.gpuid = 0
    cfg.max_iter = 90000  # the value should be set as 90000
    cfg.start_scan_idx = 0
    cfg.end_scan_idx = 1
    cwd = os.getcwd()
    cfg.pretrained_model = cwd + '/../VGG16.v2.caffemodel'
    cfg.data_root_dir = cwd + '/../../../data/'

    cfg.devkit_path = cfg.data_root_dir + 'faster_rcnn_data/v9'
    cfg.caffemodel_path = cfg.data_root_dir + 'tmp_faster_rcnn/v9/caffemodel'
    cfg.candi_dir = cfg.data_root_dir + 'tmp_faster_rcnn/v9/candidates'
    cfg.pretrained_model = cwd + '/../VGG16.v2.caffemodel'
    cfg.fold_list_path = cwd + '/../fold_list_all.npy'
    cfg.kaggle_vol_path = cfg.data_root_dir + 'kaggle/vol.hdf5'
    cfg.lidc_vol_path = cfg.data_root_dir + 'lidc/vol.hdf5'
    cfg.spie_vol_path = cfg.data_root_dir + 'spie/vol.hdf5'
    cfg.stage2_vol_path = cfg.data_root_dir + 'stage2/vol.hdf5'

    cfg.candis_thresh = 0.98
    cfg.kaggle_lungmask_path = cfg.data_root_dir + 'kaggle/lungmask.hdf5'
    cfg.lidc_lungmask_path = cfg.data_root_dir + 'lidc/lungmask.hdf5'
    cfg.spie_lungmask_path = cfg.data_root_dir + 'spie/lungmask.hdf5'
    cfg.stage2_lungmask_path = cfg.data_root_dir + 'stage2/lungmask.hdf5'
    cfg.candis_lungmask_dir = cfg.data_root_dir + 'tmp_faster_rcnn/v9/candidates_lungmask'
    cfg.pkl_dir = cfg.data_root_dir + 'faster_rcnn_candidates'

    fold_names = {1: 'fold1',
                  2: 'fold2',
                  3: 'fold3',
                  4: 'fold4',
                  5: 'fold5',
                  6: 'kaggle_beni',
                  7: 'kaggle_testset',
                  8: 'spie',
                  9: 'stage2'}

    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    print('------------------------- Training -------------------------')
    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    for fold_idx in range(1, 6):
        fold_name = fold_names[fold_idx]
        print('----------------------- training ' + fold_name + ' ---------------------')
        train_process(gpu_id=cfg.gpuid,
                      devkit_path=cfg.devkit_path,
                      fold_name=fold_name,
                      output_path=os.path.join(cfg.caffemodel_path, fold_name),
                      pretrained_model=cfg.pretrained_model,
                      train_imdb='voc_2007_trainval',
                      solver='models/pascal_voc/VGG16/faster_rcnn_end2end/solver.prototxt',
                      max_iters=cfg.max_iter
                      )

    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    print('------------------------- Testing --------------------------')
    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    for fold_idx in range(1, 9):
        fold_name = fold_names[fold_idx]
        print('----------------------- testing ' + fold_name + ' ---------------------')
        if fold_idx < 6:
            caffemodel_name = os.path.join(
                cfg.caffemodel_path, fold_name,
                'vgg16_faster_rcnn_iter_' + str(cfg.max_iter) + '.caffemodel')
        else:
            caffemodel_name = os.path.join(
                cfg.caffemodel_path, 'fold1',
                'vgg16_faster_rcnn_iter_' + str(cfg.max_iter) + '.caffemodel')
        test_process(gpu_id=cfg.gpuid,
                     caffemodel=caffemodel_name,
                     candi_dir=os.path.join(cfg.candi_dir, fold_name),
                     fold_idx=fold_idx,
                     start_scan_idx=cfg.start_scan_idx,
                     end_scan_idx=cfg.end_scan_idx,
                     fold_list_path=cfg.fold_list_path,
                     kaggle_vol_path=cfg.kaggle_vol_path,
                     lidc_vol_path=cfg.lidc_vol_path,
                     spie_vol_path=cfg.spie_vol_path,
                     stage2_vol_path=None,
                     prototxt='models/pascal_voc/VGG16/faster_rcnn_end2end/test.prototxt',
                     flag_wait=True
                     )

    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    print('------------------------- Lung Mask ------------------------')
    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    for fold_idx in range(1, 9):
        fold_name = fold_names[fold_idx]
        print('----------------------- lungmask ' + fold_name + ' ---------------------')
        cdd_filter_lungmask(thresh_r=10,
                            candis_dir=os.path.join(cfg.candi_dir, fold_name),
                            candis_lungmask_dir=os.path.join(cfg.candis_lungmask_dir, fold_name),
                            kaggle_lungmask_path=cfg.kaggle_lungmask_path,
                            lidc_lungmask_path=cfg.lidc_lungmask_path,
                            spie_lungmask_path=cfg.spie_lungmask_path,
                            stage2_lungmask_path=None,
                            kaggle_vol_path=cfg.kaggle_vol_path,
                            lidc_vol_path=cfg.lidc_vol_path,
                            spie_vol_path=cfg.spie_vol_path,
                            stage2_vol_path=None
                            )

    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    print('---------------- generate candidates list ------------------')
    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    if not os.path.exists(cfg.pkl_dir):
        os.makedirs(cfg.pkl_dir)
    gen_list_5fold(candi_dir=cfg.candis_lungmask_dir,
                   candis_thresh=cfg.candis_thresh,
                   pkl_dir=cfg.pkl_dir)
    for fold_idx in range(6, 9):
        gen_list_one_fold(candi_dir=cfg.candis_lungmask_dir,
                          fold_name=fold_names[fold_idx],
                          candis_thresh=cfg.candis_thresh,
                          pkl_dir=cfg.pkl_dir)

if __name__ == '__main__':
    main()